from JQ_Code.Hugo2046.Approximation import *
import matplotlib.pyplot as plt

def get_clf_wave(price: pd.DataFrame,
                 rate: float,
                 method: str,
                 except_dir: bool = True,
                 show_tmp: bool = False,
                 dropna: bool = True) -> pd.DataFrame:

    if except_dir:
        # 修正
        perpare_data = Pipeline([('approximation', Approximation(rate, method)),
                                 ('mask_dir_peak_valley', Mask_status_peak_valley('dir')),
                                 ('except', Except_dir('dir')),
                                 ('mask_status_peak_valley', Mask_dir_peak_valley('status'))
                                 ])
    else:
        # 普通
        perpare_data = Pipeline([('approximation', Approximation(rate, method)),
                                 ('mask_dir_peak_valley', Mask_dir_peak_valley('dir')),
                                 ('mask_status_peak_valley', Mask_status_peak_valley('dir'))])
    # 获取超参列表
    pp_params = perpare_data.get_params()

    # 数据预处理函数fit_transform()和transform()的区别 https://blog.csdn.net/Darren1921/article/details/81103277
    return perpare_data.fit_transform(price) # fit_transform()：在同一数据集上组合fit()和transform()


# 画出区别上下行的图
def plot_pivots(peak_valley_df: pd.DataFrame,
                show_dir: Union[str, List, Tuple] = 'dir',
                show_hl: bool = True,
                show_point: bool = True,
                title: str = '',
                ax=None):
    if ax is None:

        fig, ax = plt.subplots(figsize=(18, 6))

        line = peak_valley_df.plot(y='close', alpha=0.6, title=title, ax=ax)

    else:

        line = peak_valley_df.plot(y='close', alpha=0.6, title=title, ax=ax)

    if show_hl:
        peak_valley_df.plot(ax=line,
                            y='PEAK',
                            marker='o',
                            color='r',
                            mec='black')

        peak_valley_df.plot(ax=line,
                            y='VALLEY',
                            marker='o',
                            color='g',
                            mec='black')

    if show_point:
        peak_valley_df.dropna(subset=['POINT']).plot(ax=line,
                                                     y='POINT',
                                                     color='darkgray',
                                                     ls='--')
    if show_dir:
        peak_valley_df.plot(ax=line,
                            y=show_dir,
                            secondary_y=True,
                            alpha=0.3,
                            ls='--')

    return line


if __name__ == '__main__':

    from datetime import date, timedelta, datetime as dt

    startDt = str(date.today() + timedelta(days=-183))
    endDt = str(date.today() + timedelta(days=-1))

    hs300 = pd.read_csv('hs300.csv', index_col=[0], parse_dates=[0])
    # 测试方式三
    flag_frame3: pd.DataFrame = get_clf_wave(hs300,2,'c',False)
    flag_df3 = flag_frame3.loc[startDt: endDt, ['close', 'dir']]
    flag_df3.rename(columns={'dir': '方式3划分上下行'}, inplace=True)

    print('finished')